%-------------------------------------------------------------------------%    
%This program solves the model w/ and w/o a one-time 20% pay bump at age 58
%-------------------------------------------------------------------------% 

cd '/Users/dpatki/Dropbox/Research/Aging LFP/structural model/programs';
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/Disclosed output/HRS MiCDA/20210405/Tables');
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/Disclosed output/20190715/graphing');
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/data');
addpath ('/Users/dpatki/Dropbox/Research/Aging LFP/structural model/programs/lininterp'); 

%Parameters from estimated model
param = csvread('/Users/dpatki/Dropbox/Research/Aging LFP/output/x1_sobolstrt.csv');

        %=========================================================================%
        % Part 1: parameters
        %=========================================================================%
        rng(1,'twister');
        s = rng;
 	    sigma = param(1);      %EIS      
	    beta  = .97;           %discount factor 
        aT    = 81;            %terminal age
        aS    = 50;            %start age
        r     = .029;           %interest rate 
        pi    = .028;           %inflation
        V_term = 0;            %terminal value
        N_obs  = 5000;          %Number of simulated individuals
        ages = aS:1:aT-1;      %ages
        t = (1:(aT-aS))';      %time periods
        scaler = 1;            %scaler on consumption utility 
        solver_param = [sigma;beta;aT;aS;r;V_term;scaler;pi]; %collect some params
      
        %-------------------------------------------------------------------------%
        %Calculate annuity value and earnings per period from HRS
        %-------------------------------------------------------------------------%
        %1. HRS db wealth by age
        pdata = csvread('compwealth_frz_w10_22.csv',1,0);
        %col2 is age
        %col3 is avg db wealth in 2010 dollars
        %col4 is avg earnings in 2010 dollars
        dbw = pdata(aS-48+1:aT-48,2:3);
        y_raw = pdata(aS-48+1:end,4);
        %Smooth the earnings profile using cubic polynomic fit
        y_fit_param =  polyfit(ages',y_raw,3);
        y = polyval(y_fit_param,ages');
        %Earnings adjusted upward by by 20 % at age 58 (one-time shock)
        adj = ones(length(ages),1);
        adj = [1.*adj(1:8);1.2.*adj(9);1.*adj(10:end)];
        y_adj = y.*adj;
        %2. SS annuity by retirement age (hh-level)
        b_ss = csvread('ssw10.csv',1,0);
        b_ss = b_ss(aS-48+1:end,4);
        %3. SSA mortality rates by age (48+)
        pdeath = csvread('SSA_mortality_2010.csv',1,0);
        pdeath = pdeath(aS+1:end,:);
        avpdeath = mean(pdeath(:,2:3),2);
        %4. annuity factors
        afactor = zeros(length(avpdeath),1); 
        %loop over age
        for a = 1:length(afactor)
            disc = zeros(length(afactor)-a+1,1);
            disc(1) = 1;
            for i = 2:length(disc)
                disc(i) = (1-avpdeath(a+i-1))*(1/((1+r)*(1+pi)))^(i-1); 
            end

            afactor(a) = sum(disc);
        end     
        afactor = [pdeath(1:aT-aS,1) afactor(1:aT-aS)];
        %DB annuity by retirement age
        b_db = dbw(:,2)./afactor(:,2);
        %total annuity benefit
        b = b_db + b_ss;
       
        %-------------------------------------------------------------------------%
        % Normalized mortality rate
        %-------------------------------------------------------------------------%
        mortp = avpdeath(1:aT-aS);
        mortp_80 = mortp./avpdeath(aT-aS+1);
        
        %-------------------------------------------------------------------------%
        % Asset grids (A = non-pension; W = DC pension)
        %-------------------------------------------------------------------------%
        %A
        A_gridN = 20;
        split = 5;
        A_mins = [0;50000;150000;500000;1000000;2000000];
        As=zeros(A_gridN/split,split);
        for i = 1:5
            step = (A_mins(i+1)-A_mins(i))/(A_gridN/split);
            As(:,i) = A_mins(i):step:(A_mins(i+1)-step);
        end
        A = reshape(As,[A_gridN,1]);
        
        %W
        W_gridN = 20;
        split = 5;
        W_mins = [0;25000;75000;250000;500000;1500000];
        Ws=zeros(W_gridN/split,split);
        for i = 1:5
            step = (W_mins(i+1)-W_mins(i))/(W_gridN/split);
            Ws(:,i) = W_mins(i):step:(W_mins(i+1)-step);
        end
        W = reshape(Ws,[W_gridN,1]);    
        
        %-------------------------------------------------------------------------%
        % Average tax rates from nber taxsim
        %-------------------------------------------------------------------------%       
        %Control
        worktax = csvread('workinc_taxsim_a_20.csv',0,0);
        %col1 = age, col2 = earn_before_dc, col3 =fica, col4=earn_after_dc
        %col5=fiitax, col6=dc_contrib
        itau_w = worktax(:,5)./worktax(:,4); %income tax
        ptau_w = 0.5*worktax(:,3)./worktax(:,2); %worker share of FICA
        %Taxes in retirement 
        rettax = csvread('retinc_taxsim_a_20.csv',0,0);
        %col1 = fiitax, col2 = age, col3=pension distribution, col4 = SS
        itau_r = rettax(:,1)./(rettax(:,3)+rettax(:,4));
        
        %With 20% pay bump
        worktax_shock = csvread('workinc_taxsim_s58_a_20.csv',0,0);
        %col1 = age, col2 = earn_before_dc, col3 =fica, col4=earn_after_dc
        %col5=fiitax, col6=dc_contrib
        itau_w_shock = worktax_shock(:,5)./worktax_shock(:,4); %income tax
        ptau_w_shock = 0.5*worktax_shock(:,3)./worktax_shock(:,2); %worker share of FICA  
             
        %Collect parameters for model solution
        %Basline rates
        taxf = cell(3,1);
        taxf{1,1}=itau_w;
        taxf{2,1}=ptau_w;
        taxf{3,1}=itau_r;
        
        %Rates w/ pay bump
        taxf_shock = cell(3,1);
        taxf_shock{1,1}=itau_w_shock;
        taxf_shock{2,1}=ptau_w_shock;
        taxf_shock{3,1}=itau_r;        
        
        %-------------------------------------------------------------------------%
        %Disutility of work process 
        %-------------------------------------------------------------------------%
    	mu      = 0;  %mean normal underlying f
        rho     = 1/(1+exp(param(2)));  %persistence of random process (random walk)
    	sigma_v = exp(param(3)); %sd of normal underlying f
    	phi     = param(4); %age slope
        gamma   = param(5); %constant
    	f_N     = 6; %grid size

        %=========================================================================%
        % Part 2: solve model for value and policy function 
        % under basline condition (no pay bump)
        %=========================================================================%

        %Cell array to store value and policy functions for each age
        optfunc = cell(length(t),8);

        %-------------------------------------------------------------------------%
        %Solve age T-1 policy and value functions outside the loop
        %-------------------------------------------------------------------------%	
        %Set last period=true
        last=1;
        a=0;
        [V_w,A_polw,W_polw]=vwork_5(A,W,last,a,y,optfunc,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxf,mortp_80);
        [V_r,A_polr,W_polr]=vret_4(A,W,last,a,b_db,b_ss,optfunc,solver_param,taxf,mortp_80);
	
        
        % Outer value function and labor supply choice
        V = zeros(size(V_w)); 
        L = zeros(size(V_w)); 
       
        for i = 1:length(A)
            for j = 1:length(W)
                for k = 1:f_N
                    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
                    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
                end
            end
        end

        %Store value and policy functions
        optfunc(1,:) = {V_w;A_polw;W_polw;V_r;A_polr;W_polr;V;L};
        
        %-------------------------------------------------------------------------%
        %Solve policy and value functions for remaning periods
        %-------------------------------------------------------------------------%      
        %Set last period = false
        last=0;
        
        %Loop through ages
        for a = 1:length(t)-1 
    

               [V_w,A_polw,W_polw]=vwork_5(A,W,last,a,y,optfunc,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxf,mortp_80);
               [V_r,A_polr,W_polr]=vret_4(A,W,last,a,b_db,b_ss,optfunc,solver_param,taxf,mortp_80);


		% Outer value function and labor supply choice
		V = zeros(size(V_w)); 
		L = zeros(size(V_w)); 

		for i = 1:length(A)
		    for j = 1:length(W)
			for k = 1:f_N
			    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
			    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
			end
		    end
		end

		%Store value and policy functions
        optfunc(a+1,:) = {V_w;A_polw;W_polw;V_r;A_polr;W_polr;V;L};

        end
        
        %=========================================================================%
        % Part 3: solve model for value and policy function 
        % under 20 %, one-time pay bump condition
        %=========================================================================%

        %Cell array to store value and policy functions for each age
        optfunc_shock = cell(length(t),8);

        %-------------------------------------------------------------------------%
        %Solve age T-1 policy and value functions outside the loop
        %-------------------------------------------------------------------------%	
        %Set last period=true
        last=1;
        a=0;
        [V_w,A_polw,W_polw]=vwork_5(A,W,last,a,y_adj,optfunc_shock,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxf_shock,mortp_80);
        [V_r,A_polr,W_polr]=vret_4(A,W,last,a,b_db,b_ss,optfunc_shock,solver_param,taxf_shock,mortp_80);
	
        
        % Outer value function and labor supply choice
        V = zeros(size(V_w)); 
        L = zeros(size(V_w)); 
       
        for i = 1:length(A)
            for j = 1:length(W)
                for k = 1:f_N
                    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
                    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
                end
            end
        end

        %Store value and policy functions
        optfunc_shock(1,:) = {V_w;A_polw;W_polw;V_r;A_polr;W_polr;V;L};
        
        %-------------------------------------------------------------------------%
        %Solve policy and value functions for remaning periods
        %-------------------------------------------------------------------------%      
        %Set last period = false
        last=0;
        
        %Loop through ages
        for a = 1:length(t)-1 
    

               [V_w,A_polw,W_polw]=vwork_5(A,W,last,a,y_adj,optfunc_shock,solver_param,...
                    f_N,mu,rho,sigma_v,phi,gamma,taxf_shock,mortp_80);
               [V_r,A_polr,W_polr]=vret_4(A,W,last,a,b_db,b_ss,optfunc_shock,solver_param,taxf_shock,mortp_80);


		% Outer value function and labor supply choice
		V = zeros(size(V_w)); 
		L = zeros(size(V_w)); 

		for i = 1:length(A)
		    for j = 1:length(W)
			for k = 1:f_N
			    V(i,j,k) = max(V_w(i,j,k),V_r{aT-aS-a}(i,j));
			    L(i,j,k) = (V_w(i,j,k)>V_r{aT-aS-a}(i,j));
			end
		    end
		end

		%Store value and policy functions
        optfunc_shock(a+1,:) = {V_w;A_polw;W_polw;V_r;A_polr;W_polr;V;L};

        end
        
      
        %=========================================================================%
        %Part 4: Simulate shocks and solve for labor supply and saving
        %=========================================================================%
        
        %Cell arrays to hold simulated paths 
        LsimC = cell(1,1);
        AsimC = cell(1,1);
        WsimC = cell(1,1);
        LsimT = cell(1,1);
        AsimT = cell(1,1);
        WsimT = cell(1,1);
      
        %Nodes of discrete AR(1) process -- need max and min of these to
        %assign values out of bounds in the simulation
        [f,~] = rouwenhorst(f_N,mu,rho,sigma_v);
        hor = 12; %horizon (15 years post)
        
        %Simulated observations 
        obs = N_obs;
        
        %Disutility shock process
        %simulate AR(1) shocks. 1000 period burn-in... 
        f_sim0 = zeros(1000+(hor+6),N_obs);
        rng(s);
        v = normrnd(0,sigma_v,[length(f_sim0),N_obs]);
        for i = 2:length(f_sim0)
            f_sim0(i,:) = (1-rho)*mu+rho*f_sim0(i-1,:)+v(i,:);
        end
        %Trim away burn-in periods
        f_sim0 = f_sim0(end-(hor+6-1):end,:);
        D{1,1} = f_sim0;
       
        %Initial assets (age 53) --moments from HRS        
        A0 = cell(1,1);
        W0 = cell(1,1);
        A_params = csvread('A_params.csv',1,0);
        J_params = csvread('J_params.csv',1,0);
        rng(s);    
        %Number of individuals with 0 DC wealth
        frac0obs = round(A_params(3,4)*N_obs); 
        %Non-DC wealth for individuals with 0 DC wealth
        A_temp0 = lognrnd(A_params(3,2),A_params(3,3),[frac0obs,1]);
        %DC wealth for individuals with no DC balances = 0
        W_temp0 = zeros(frac0obs,1);
        %Joint wealth distribution for individuals with >0 DC wealth
        mvnmu = J_params(3,2:3);
        mvnsig = [(J_params(3,4))^2 J_params(3,6);J_params(3,6) (J_params(3,5))^2];
        J_wealth = exp(mvnrnd(mvnmu,mvnsig,(N_obs-frac0obs)));
        W_temp1 = J_wealth(:,1);
        A_temp1 = J_wealth(:,2);
        A0{1,1} = [A_temp0' A_temp1'];
        W0{1,1} = [W_temp0' W_temp1'];       

        %use policy functions to obtain optimal labor supply and asset
        %accumulation for each simulated individual in each period 
        %-------------------------------------------------------------------------%
        %Baseline simulation
        %-------------------------------------------------------------------------%
        
        p=3; %Age 53 in period -5

        L_sim = zeros(hor+6,N_obs);
        A_sim = zeros(hor+6,N_obs);
        W_sim = zeros(hor+6,N_obs);
        A_sim(1,:) = A0{1,1};
        W_sim(1,:) = W0{1,1};

        %Set disutility distribution
        f_sim = D{1,1};

        %Employment choice for t = -5 outside the loop (t = period
        %relative to shock) a = 1 when t = -5
        for i = 1:N_obs
            a=1;    
            Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
            Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
            Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
            Afunc_r  = Afunc_r1{p+a};
            Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
            Wfunc_r  = Wfunc_r1{p+a};
            Lfunc = optfunc{(aT-aS)-(p+a-1),8};
            %Put extreme draws of A,W,f onto the grid 
            %A
            if A_sim(a,i)>max(A)
                A_sim(a,i) = max(A);
            end
            %W
            if W_sim(a,i)>max(W)
                W_sim(a,i) = max(W);
            end	
            %f                
            if f_sim(a,i)<min(f) 
                f_sim(a,i) = min(f);
            elseif f_sim(a,i)>max(f)
                f_sim(a,i) = max(f);
            end      

            %Simulated values 
            L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));

            if L_sim(a,i) == 1
                A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
            elseif L_sim(a,i) == 0 
                A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
            end 
        end

        for a = 2:hor+5 
           for i = 1:N_obs

                Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end                    
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                end

            end
        end

        %Final period employment choice (A,W already determined)
        a = hor+6;
        for i = 1:N_obs

                Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                                     
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                end

        end

        LsimC{1,1} = L_sim;
        AsimC{1,1} = A_sim;
        WsimC{1,1} = W_sim;
    
        %-------------------------------------------------------------------------%
        %20% earnings shock at age 58
        %-------------------------------------------------------------------------%
    
        p=3;  %Age 53 in period -5
        L_sim = zeros(hor+6,N_obs);
        A_sim = zeros(hor+6,N_obs);
        W_sim = zeros(hor+6,N_obs);
        A_sim(1,:) = A0{1,1};
        W_sim(1,:) = W0{1,1};

        %Set disutility distribution
        f_sim = D{1,1};

        %Employment choice for t = -5 outside the loop (t = period
        %relative to shock) a = 1 when t = -5
        for i = 1:N_obs
            a=1;    
            Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
            Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
            Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
            Afunc_r  = Afunc_r1{p+a};
            Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
            Wfunc_r  = Wfunc_r1{p+a};
            Lfunc = optfunc{(aT-aS)-(p+a-1),8};
            %Put extreme draws of A,W,f onto the grid 
            %A
            if A_sim(a,i)>max(A)
                A_sim(a,i) = max(A);
            end
            %W
            if W_sim(a,i)>max(W)
                W_sim(a,i) = max(W);
            end	
            %f                
            if f_sim(a,i)<min(f) 
                f_sim(a,i) = min(f);
            elseif f_sim(a,i)>max(f)
                f_sim(a,i) = max(f);
            end      

            %Simulated values 
            L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));
            if L_sim(a,i) == 1
                A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
            elseif L_sim(a,i) == 0 
                A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
            end    
        end

        for a = 2:5 
           for i = 1:N_obs

                Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end                    
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                end

            end
        end
        
       for a = 6:hor+5 
           for i = 1:N_obs

                Afunc_w = optfunc_shock{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc_shock{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc_shock{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc_shock{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc_shock{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end                    
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                end

            end
       end
        
        %Final period employment choice (A,W already determined)
        a = hor+6;
        for i = 1:N_obs

                Afunc_w = optfunc_shock{(aT-aS)-(p+a-1),2};
                Wfunc_w = optfunc_shock{(aT-aS)-(p+a-1),3};      
                Afunc_r1 = optfunc_shock{(aT-aS)-(p+a-1),5};
                Afunc_r  = Afunc_r1{p+a};
                Wfunc_r1 = optfunc_shock{(aT-aS)-(p+a-1),6};  
                Wfunc_r  = Wfunc_r1{p+a};                  
                Lfunc = optfunc_shock{(aT-aS)-(p+a-1),8};

                %Put extreme draws of A,W,f onto the grid 
                %A
                if A_sim(a,i)>max(A)
                    A_sim(a,i) = max(A);
                end
                %W
                if W_sim(a,i)>max(W)
                    W_sim(a,i) = max(W);
                end	
                %f                
                if f_sim(a,i)<min(f) 
                    f_sim(a,i) = min(f);
                elseif f_sim(a,i)>max(f)
                    f_sim(a,i) = max(f);
                end        

                %Enforce self-absorbing retirement (value functions impose it)
                if  L_sim(a-1,i) == 1                 
                    L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                                  
                elseif L_sim(a-1,i) == 0
                    L_sim(a,i) = 0;
                end

        end

        LsimT{1,1} = L_sim;
        AsimT{1,1} = A_sim;
        WsimT{1,1} = W_sim;

    %-------------------------------------------------------------------------%
    % Extract simulated data
    %-------------------------------------------------------------------------%

    %Extract control group       
    L_C2 = LsimC{1,1};
    A_C2 = AsimC{1,1};
    W_C2 = WsimC{1,1};

    %Extract treated group data        
    L_T2 = LsimT{1,1};
    A_T2 = AsimT{1,1};
    W_T2 = WsimT{1,1};    
    
%-------------------------------------------------------------------------%    
%Export data for analysis
%-------------------------------------------------------------------------% 
%Create some inputs
aT    = 81;            %terminal age
aS    = 50;            %start age
pdata = csvread('compwealth_frz_w10_22.csv',1,0);
ages = aS:1:aT-1;      %ages
y_raw = pdata(aS-48+1:end,4); %raw earnings profile
y_fit_param =  polyfit(ages',y_raw,3);
y = polyval(y_fit_param,ages');  %smoothed earnings profile
ernstart = 4; %start at age 53
%SSA mortality rates by age (48+)
pdeath = csvread('SSA_mortality_2010.csv',1,0);
pdeath = pdeath(aS+1:end,:);
avpdeath = mean(pdeath(:,2:3),2);
mortp = avpdeath(1:aT-aS);
mortp_80 = mortp./avpdeath(aT-aS+1);
    
%------------------------------
%Baseline group
%------------------------------

%Extract simulated data 
L = L_C2;
A = A_C2;
W = W_C2;
ages = (53:1:70)';
g = D{1,1}+gamma+phi.*ages;

%Remove non-workers in -1 period
indx_1 = find(L(5,:)==0);
L(:,indx_1)=[];
A(:,indx_1)=[];
W(:,indx_1)=[];
g(:,indx_1)=[];

%Create id's
ids = zeros(size(L));
for i = 1:length(L)
    ids(:,i) = i;
end

ids_vec = reshape(ids,[18*length(L),1]);
LT = reshape(L,[18*length(L),1]);
AT = reshape(A,[18*length(L),1]);
WT = reshape(W,[18*length(L),1]);
gT = reshape(g,[18*length(L),1]);

%Earnings 
ern = y(ernstart:ernstart+17);
ernrep = repmat(ern,[length(L),1]);
%SS benefits
ssben = b_ss(ernstart:ernstart+17);
ssbenrep = repmat(ssben,[length(L),1]);
%DB benefits
B_db = b_db(ernstart:ernstart+17);
B_dbrep = repmat(B_db,[length(L),1]);

%Probability of death
mort = mortp_80(ernstart:ernstart+17);
mortrep = repmat(mort,[length(L),1]);

Cdata = [ids_vec gT LT AT WT ernrep B_dbrep ssbenrep mortrep];

%------------------------------
%20% pay bump group
%------------------------------
 
%Extract simulated data 
L = L_T2;
A = A_T2;
W = W_T2;
ages = (53:1:70)';
g = D{1,1}+gamma+phi.*ages;

%Remove non-workers in -1 period
indx_1 = find(L(5,:)==0);
L(:,indx_1)=[];
A(:,indx_1)=[];
W(:,indx_1)=[];
g(:,indx_1)=[];

%Create id's
ids = zeros(size(L));
for i = 1:length(L)
    ids(:,i) = i;
end 

ids_vec = reshape(ids,[18*length(L),1]);
LT = reshape(L,[18*length(L),1]);
AT = reshape(A,[18*length(L),1]);
WT = reshape(W,[18*length(L),1]);
gT = reshape(g,[18*length(L),1]);
    
%Earnings 
ern = y_adj(ernstart:ernstart+17);
ernrep = repmat(ern,[length(L),1]);
%SS benefits
ssben = b_ss(ernstart:ernstart+17);
ssbenrep = repmat(ssben,[length(L),1]);
%DB benefits
B_db = b_db(ernstart:ernstart+17);
B_dbrep = repmat(B_db,[length(L),1]);

%Probability of death
mort = mortp_80(ernstart:ernstart+17);
mortrep = repmat(mort,[length(L),1]);

Tdata = [ids_vec gT LT AT WT ernrep B_dbrep ssbenrep mortrep];

%Stack control and treated data
trt_flag = ones(length(Tdata),1);
Cdata = [Cdata (1-trt_flag)];
Tdata = [Tdata trt_flag];
data  = [Cdata;Tdata];

csvwrite('/Users/dpatki/Dropbox/Research/Aging LFP/data/sim_annfactor.csv',data);

     